# -*- encoding: utf8 -*-
import numpy as np
import matplotlib.pyplot as plt

def nextpow2(n):
    """
    子函数 nextpow2(): 求最接近数据长度的2的整数次方

    Eg:
    nextpow2(2) = 1
    nextpow2(1025) = 11
    nextpow2(2**20+1) = 21
    """
    return np.ceil(np.log2(np.abs(n))).astype('long')

def main_program():
    """主程序: 绘制原始信号波形 + 傅里叶频域分析图"""

    # 输出图片画框尺寸设置
    fig, ax = plt.subplots(2, 1, figsize=(18, 6))
    fig.tight_layout(h_pad=6, w_pad=2)
    plt.rcParams['font.sans-serif'] = ['KaiTi'] # 楷体
    plt.rcParams['axes.unicode_minus'] = False # 避免坐标轴的无法正常显示负号

    # ------------------
    # 0.参数
    Fs = 10000  # 采样率=10kHz
    T  = 1.0/Fs # 采样间隔0.1毫秒
    L  = 1500   # 信号长度（1500）
    t  = T * np.arange(0, (L-1)) # 步进间隔T=0.1ms, 共1500个采样点
    HZ = 800    # 信号频率Hz

    # ------------------
    # 一、原始数据的波形
    w = HZ * 2 * np.pi
    y = 12 * np.sin(w * t)
    ax1 = plt.subplot(2, 1, 1)
    ax1.plot(t, y, '-', linewidth=1, label="原始数据波形") # 绘图并设置线条样式
    ax1.set_title("输入(800Hz正弦波,采样率=10kHz)")
    ax1.set_xlabel("t/s")
    ax1.set_ylabel("y/V")
    ax1.legend()
    ax1.grid(axis = 'y', linestyle = '-')

    # ----------------
    # 二、取接近2^n方便计算,1024 , 2048
    FFTNUM = 2 ** nextpow2(1500) #求得最接近采样点的2^n,由于上面是1500点,那么最近的就是2048点。
    #FFTNUM = L
    Y = np.fft.fft(y, FFTNUM) / L

    # ------------------
    # 三、 0~5000hz,均分为751份（或1025份）
    f = Fs/2 * np.linspace(0, 1, int(FFTNUM/2 + 1)) # 频率轴。这里只显示Fs/2前半部分,另一半是对称的。
    yy = 2 * np.abs(Y[:int(FFTNUM/2 + 1)])
    ax2 = plt.subplot(2, 1, 2)
    ax2.plot(f, yy, color="red", label="傅里叶变换")
    ax2.set_title("频域分析")
    ax2.set_xlabel("f/Hz")
    ax2.set_ylabel("yy/V")
    ax2.legend()
    ax2.grid(axis = 'y', linestyle = '-')

    # -------------------------
    # 四、输出调试参数
    hz = w / (2 * np.pi)
    MaxData = np.max(yy)
    print("信号频率 hz = %.2f" %(hz))
    print("max = %.2f" %(MaxData))

# Python脚本入口
if __name__ == '__main__':
    # 调用主程序
    main_program()
    fig_manager = plt.get_current_fig_manager()
    (width, height) = fig_manager.window.maxsize()
    fig_manager.resize(width - 100, height - 200)
    fig_manager.window.state('zoomed') # 全屏最大化

    # 保存图片
    plt.savefig("./output.png", bbox_inches="tight") # savefig() 必须要放在plt.show()前面，show()之后会自动释放图表内存
    plt.draw()
    plt.show()
    # exit(0)
